package anthropic

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/mylxsw/go-utils/array"
	"io"
	"net/http"
	"strings"
)

type Anthropic struct {
	apiKey    string
	serverURL string

	client *http.Client
}

type Model string

const (
	ModelClaudeInstant Model = "claude-instant-1"
	ModelClaude2       Model = "claude-2"
	ModelClaude3Opus   Model = "claude-3-opus"
	ModelClaude3Sonnet Model = "claude-3-sonnet"
	ModelClaude3Haiku  Model = "claude-3-haiku"
)

func New(serverURL, apiKey string, client *http.Client) *Anthropic {
	if serverURL == "" {
		serverURL = "https://api.anthropic.com"
	}

	if client == nil {
		client = http.DefaultClient
	}

	return &Anthropic{apiKey: apiKey, serverURL: serverURL, client: client}
}

type MessageRequest struct {
	// Model The model that will complete your prompt.
	Model Model `json:"model"`
	// Messages The messages that you want Claude to complete.
	Messages []Message `json:"messages"`
	// System System prompt
	System string `json:"system,omitempty"`
	// MaxTokens The maximum number of tokens to generate before stopping.
	// Note that our models may stop before reaching this maximum.
	// This parameter only specifies the absolute maximum number of tokens to generate.
	MaxTokens int `json:"max_tokens,omitempty"`
	// Stream Whether to incrementally stream the response using server-sent events.
	Stream bool `json:"stream,omitempty"`
	// Temperature Amount of randomness injected into the response.
	// Defaults to 1.0. Ranges from 0.0 to 1.0.
	// Use temperature closer to 0.0 for analytical / multiple choice, and closer to 1.0 for creative and generative tasks.
	Temperature float64 `json:"temperature,omitempty"`
	// TopP Use nucleus sampling.
	// In nucleus sampling, we compute the cumulative distribution over all the options for each subsequent token in
	// decreasing probability order and cut it off once it reaches a particular probability specified by top_p.
	// You should either alter temperature or top_p, but not both.
	// Recommended for advanced use cases only. You usually only need to use temperature.
	TopP float64 `json:"top_p,omitempty"`
	// TopK only sample from the top K options for each subsequent token.
	// Used to remove "long tail" low probability responses. Learn more technical details here.
	// Recommended for advanced use cases only. You usually only need to use temperature.
	TopK int `json:"top_k,omitempty"`

	// Thinking 扩展思维
	Thinking *Thinking `json:"thinking,omitempty"`
}

type Thinking struct {
	Type         string `json:"type,omitempty"`
	BudgetTokens int    `json:"budget_tokens,omitempty"`
}

type Message struct {
	// Role The role of the message.
	Role string `json:"role"`
	// Content The content of the message.
	Content []MessageContent `json:"content"`
}

func NewTextMessage(role string, text string) Message {
	return Message{
		Role:    role,
		Content: []MessageContent{{Type: "text", Text: text}},
	}
}

type MessageContent struct {
	// Type The type of the message, support "text", "image"
	Type string `json:"type"`
	// Text The text of the message. Required if type is "text".
	Text string `json:"text,omitempty"`
	// Source The source of the image. Required if type is "image".
	Source *ImageSource `json:"source,omitempty"`
}

func NewImageSource(mediaType, data string) *ImageSource {
	return &ImageSource{
		Type:      "base64",
		MediaType: mediaType,
		Data:      data,
	}
}

type ImageSource struct {
	// Type The type of the image source, only support "base64"
	Type string `json:"type"`
	// Data The base64 encoded image data, such as image/jpeg, image/png, image/gif, image/webp.
	MediaType string `json:"media_type"`
	// Data The base64 encoded image data.
	Data string `json:"data"`
}

type MessageResponse struct {
	// ID Unique object identifier.
	// The format and length of IDs may change over time.
	ID string `json:"id,omitempty"`
	// Type Object type.
	// For Messages, this is always "message".
	Type string `json:"type,omitempty"`
	// Role Conversational role of the generated message.
	//This will always be "assistant".
	Role string `json:"role,omitempty"`
	// Content generated by the model.
	// This is an array of content blocks, each of which has a type that determines its shape.
	// Currently, the only type in responses is "text"
	Content []MessageResponseContent `json:"content,omitempty"`
	// Model The model that handled the request.
	Model string `json:"model,omitempty"`
	// StopReason The reason that we stopped.
	// This may be one the following values:
	// - "end_turn": the model reached a natural stopping point
	// - "max_tokens": we exceeded the requested max_tokens or the model's maximum
	// - "stop_sequence": one of your provided custom stop_sequences was generated
	// Note that these values are different than those in /v1/complete, where end_turn and stop_sequence were not differentiated.
	// In non-streaming mode this value is always non-null. In streaming mode, it is null in the message_start event and non-null otherwise.
	StopReason string `json:"stop_reason,omitempty"`
	// StopSequence Which custom stop sequence was generated, if any.
	// This value will be a non-null string if one of your custom stop sequences was generated.
	StopSequence string `json:"stop_sequence,omitempty"`
	// Usage Billing and rate-limit usage.
	// Anthropic's API bills and rate-limits by token counts, as tokens represent the underlying cost to our systems.
	// Under the hood, the API transforms requests into a format suitable for the model.
	// The model's output then goes through a parsing stage before becoming an API response.
	// As a result, the token counts in usage will not match one-to-one with the exact visible content of an API request or response.
	// For example, output_tokens will be non-zero, even for an empty string response from Claude.
	Usage *Usage `json:"usage,omitempty"`

	// Error 错误信息
	Error *ResponseError `json:"error,omitempty"`
}

func (resp MessageResponse) Text() string {
	var res string
	for _, content := range resp.Content {
		res += content.Text
	}

	return res
}

type Usage struct {
	InputTokens  int `json:"input_tokens,omitempty"`
	OutputTokens int `json:"output_tokens,omitempty"`
}

type MessageResponseContent struct {
	Type string `json:"type,omitempty"`
	Text string `json:"text,omitempty"`
}

func (ai *Anthropic) Chat(ctx context.Context, req MessageRequest) (*MessageResponse, error) {
	req.Stream = false
	if req.MaxTokens <= 0 {
		req.MaxTokens = 20000
	}

	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("marshal request failed: %s", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", strings.TrimRight(ai.serverURL, "/")+"/v1/messages", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create http request failed: %s", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("anthropic-version", "2023-06-01")
	httpReq.Header.Set("x-api-key", ai.apiKey)

	httpResp, err := ai.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("chat failed: %s", err)
	}

	defer httpResp.Body.Close()

	if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusBadRequest {
		data, _ := io.ReadAll(httpResp.Body)
		return nil, fmt.Errorf("chat failed [%s]: %s", httpResp.Status, string(data))
	}

	var chatResp MessageResponse
	if err := json.NewDecoder(httpResp.Body).Decode(&chatResp); err != nil {
		return nil, fmt.Errorf("decode response failed: %s", err)
	}

	return &chatResp, nil
}

type MessageStreamResponse struct {
	// Type Each event will use an SSE event name (e.g. event: message_stop), and include the matching event type in its data.
	// Each stream uses the following event flow:
	// - message_start: contains a Message object with empty content.
	// A series of content blocks, each of which have a content_block_start,
	// one or more content_block_delta events, and a content_block_stop event.
	// Each content block will have an index that corresponds to its index in the final Message content array.
	// One or more message_delta events, indicating top-level changes to the final Message object.
	// A final message_stop event.
	Type  string        `json:"type"`
	Index int           `json:"index,omitempty"`
	Delta *MessageDelta `json:"delta,omitempty"`
	// Error 错误信息
	Error *ResponseError `json:"error,omitempty"`
}

func (res MessageStreamResponse) Text() string {
	if res.Delta != nil {
		return res.Delta.Text
	}

	return ""
}

func (res MessageStreamResponse) Thinking() string {
	if res.Delta != nil {
		return res.Delta.Thinking
	}

	return ""
}

type MessageDelta struct {
	Type         string `json:"type,omitempty"`
	Text         string `json:"text,omitempty"`
	Thinking     string `json:"thinking,omitempty"`
	StopReason   string `json:"stop_reason,omitempty"`
	StopSequence string `json:"stop_sequence,omitempty"`
	Usage        *Usage `json:"usage,omitempty"`
}

func (ai *Anthropic) ChatStream(ctx context.Context, req MessageRequest) (<-chan MessageStreamResponse, error) {
	req.Stream = true
	body, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", strings.TrimRight(ai.serverURL, "/")+"/v1/messages", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("anthropic-version", "2023-06-01")
	httpReq.Header.Set("anthropic-beta", "messages-2023-12-15")
	httpReq.Header.Set("x-api-key", ai.apiKey)

	httpReq.Header.Set("Accept", "text/event-stream")
	httpReq.Header.Set("Cache-Control", "no-cache")
	httpReq.Header.Set("Connection", "keep-alive")

	httpResp, err := ai.client.Do(httpReq)
	if err != nil {
		return nil, err
	}

	if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusBadRequest {
		data, _ := io.ReadAll(httpResp.Body)
		_ = httpResp.Body.Close()

		return nil, fmt.Errorf("chat failed [%s]: %s", httpResp.Status, string(data))
	}

	res := make(chan MessageStreamResponse)
	go func() {
		defer func() {
			_ = httpResp.Body.Close()
			close(res)
		}()

		reader := bufio.NewReader(httpResp.Body)
		for {
			data, err := reader.ReadBytes('\n')
			if err != nil {
				if err == io.EOF {
					return
				}

				select {
				case <-ctx.Done():
				case res <- MessageStreamResponse{Error: &ResponseError{Type: "read_error", Message: fmt.Sprintf("read response failed: %v", err)}}:
				}
				return
			}

			dataStr := strings.TrimSpace(string(data))
			if dataStr == "" {
				continue
			}

			if !strings.HasPrefix(dataStr, "data: ") {
				continue
			}

			var chatResponse MessageStreamResponse
			if err := json.Unmarshal([]byte(dataStr[6:]), &chatResponse); err != nil {
				select {
				case <-ctx.Done():
				case res <- MessageStreamResponse{Error: &ResponseError{Type: "decode_error", Message: fmt.Sprintf("decode response failed: %v", err)}}:
				}
				return
			}

			if chatResponse.Type == "message_stop" {
				return
			}

			if array.In(chatResponse.Type, []string{"content_block_delta", "message_delta"}) && chatResponse.Delta != nil {
				select {
				case <-ctx.Done():
					return
				case res <- chatResponse:
					if chatResponse.Delta.StopReason != "" {
						return
					}
				}
			}
		}
	}()

	return res, nil
}

type ResponseError struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}
